import numpy as np
import open3d as o3d


def rot_diff_rad(r0, r1, sym_axis=None):
    if sym_axis == "x":
        x1, x2 = r0[..., 0], r1[..., 0]  # [3]
        diff = np.sum(x1 * x2, axis=-1)  # [1]
        diff = np.clip(diff, a_min=-1.0, a_max=1.0)
        return np.arccos(diff)
    elif sym_axis == "y":
        y1, y2 = r0[..., 1], r1[..., 1]  # [3]
        diff = np.sum(y1 * y2, axis=-1)  # [1]
        diff = np.clip(diff, a_min=-1.0, a_max=1.0)
        return np.arccos(diff)
    elif sym_axis == "z":
        z1, z2 = r0[..., 2], r1[..., 2]  # [3]
        diff = np.sum(z1 * z2, axis=-1)  # [1]
        diff = np.clip(diff, a_min=-1.0, a_max=1.0)
        return np.arccos(diff)
    else:
        mat_diff = np.matmul(r0, r1.swapaxes(-1, -2))
        diff = mat_diff[0, 0] + mat_diff[1, 1] + mat_diff[2, 2]
        diff = (diff - 1) / 2.0
        diff = np.clip(diff, a_min=-1.0, a_max=1.0)
        return np.arccos(diff)


def get_pose_diff(pose0, pose1, sym_axis=None):
    t0 = pose0["translation"].reshape(3, 1)
    r0 = pose0["rotation"]
    t1 = pose1["translation"].reshape(3, 1)
    r1 = pose1["rotation"]
    tdiff = np.linalg.norm(t1 - t0, ord=2)
    rdiff = rot_diff_rad(r0, r1, sym_axis=sym_axis) / np.pi * 180.0
    return tdiff, rdiff


if __name__ == "__main__":
    extr_old = o3d.io.read_pinhole_camera_trajectory("output_old_0.log")
    extr_new = o3d.io.read_pinhole_camera_trajectory("output_new_0.log")
    cam_poses_old = []
    for i in range(len(extr_old.parameters)):
        T = np.linalg.inv(extr_old.parameters[i].extrinsic)
        cam_poses_old.append({
            "rotation": T[:3, :3],
            "translation": T[:3, 3].reshape(3, 1),
        })
    cam_poses_new = []
    for i in range(len(extr_new.parameters)):
        T = np.linalg.inv(extr_new.parameters[i].extrinsic)
        cam_poses_new.append({
            "rotation": T[:3, :3],
            "translation": T[:3, 3].reshape(3, 1),
        })
    
    assert len(cam_poses_old) == len(cam_poses_new)

    tdiffs, rdiffs = [], []
    for (cam_pose_old, cam_pose_new) in zip(cam_poses_old, cam_poses_new):
        tdiff, rdiff = get_pose_diff(cam_pose_old, cam_pose_new, sym_axis=None)
        tdiffs.append(tdiff)
        rdiffs.append(rdiff)
    tdiffs = np.float32(tdiffs)
    rdiffs = np.float32(rdiffs)
    print(np.mean(tdiffs), np.std(tdiffs), np.mean(rdiffs), np.std(rdiffs))
